import cv2
import numpy as np
import os
import csv
import sys
from datetime import datetime

# 设置常量
FRAMES_FOLDER = "frames"  # 输入文件夹，包含多个子文件夹
OUTPUT_FOLDER = "visual_analysis_results"  # 输出文件夹

class SimpleImageAnalyzer:
    """简化版图像帧分析器"""
    
    def __init__(self):
        print("初始化图像分析器...")
        print("使用基本图像处理方法进行分析")
        print("分析器初始化完成！")
    
    def analyze_image_folder(self, folder_path):
        """分析单个子文件夹中的图片帧"""
        folder_name = os.path.basename(folder_path)
        print(f"开始分析子文件夹 '{folder_name}' 中的图片帧...")
        
        # 获取该文件夹的所有图片
        image_paths = [os.path.join(folder_path, f) 
                      for f in sorted(os.listdir(folder_path)) 
                      if f.lower().endswith(('.jpg', '.jpeg', '.png'))]
        
        if not image_paths:
            print(f"错误: '{folder_path}' 中没有发现图片文件")
            return None
        
        print(f"找到 {len(image_paths)} 个图片帧")
        
        # 初始化结果
        frame_results = []
        emotion_data = []
        face_detection_count = 0
        
        # 处理每一帧
        for i, image_path in enumerate(image_paths):
            if i % 10 == 0:
                print(f"处理帧 {i+1}/{len(image_paths)}...")
                
            frame = cv2.imread(image_path)
            if frame is None:
                print(f"警告: 无法读取图片 '{image_path}'")
                continue
                
            gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
            frame_result = {
                'path': image_path,
                'eye_contact': 0.0,
                'expression': 0.0,
                'posture': 0.0,
                'emotion': 'unknown'
            }
            
            # 人脸检测和特征分析
            h, w = gray.shape
            regions = {
                'upper': gray[0:h//3, :],
                'middle': gray[h//3:2*h//3, :],
                'lower': gray[2*h//3:h, :],
                'center': gray[h//4:3*h//4, w//4:3*w//4]
            }
            
            variances = {k: np.var(v) for k, v in regions.items()}
            brightness = {k: np.mean(v) for k, v in regions.items()}
            
            overall_brightness = np.mean(gray)
            overall_contrast = np.std(gray)
            
            face_detected = (variances['center'] > 200 and 
                           overall_brightness > 30 and
                           overall_contrast > 40)
            
            if face_detected:
                face_detection_count += 1
                eye_score = min(1.0, variances['upper'] / 2000)
                frame_result['eye_contact'] = float(eye_score)
                
                expr_score = min(1.0, variances['lower'] / 1500)
                frame_result['expression'] = float(expr_score)
                
                center_offset = abs(brightness['center'] - overall_brightness) / overall_brightness
                posture_score = 1.0 - min(1.0, center_offset * 2)
                frame_result['posture'] = float(posture_score)
                
                if expr_score > 0.6:
                    frame_result['emotion'] = 'happy'
                elif expr_score > 0.4:
                    frame_result['emotion'] = 'neutral'
                else:
                    frame_result['emotion'] = 'serious'
                
                emotion_data.append(frame_result['emotion'])
                
                if i > 0 and 'prev_gray' in locals():
                    if prev_gray.shape == gray.shape:
                        frame_result['face_dynamics'] = np.mean(np.abs(gray.astype(float) - prev_gray.astype(float)))
                    else:
                        try:
                            resized_prev = cv2.resize(prev_gray, (gray.shape[1], gray.shape[0]))
                            frame_result['face_dynamics'] = np.mean(np.abs(gray.astype(float) - resized_prev.astype(float)))
                        except Exception as e:
                            print(f"警告: 无法比较不同尺寸的帧: {e}")
                            frame_result['face_dynamics'] = 0.0
                else:
                    frame_result['face_dynamics'] = 0.0
                
                prev_gray = gray.copy()
            else:
                frame_result['eye_contact'] = 0.3
                frame_result['expression'] = 0.3
                frame_result['posture'] = 0.5
                frame_result['emotion'] = 'unknown'
                frame_result['face_dynamics'] = 0.0
                emotion_data.append('unknown')
            
            frame_results.append(frame_result)
        
        # 汇总结果
        total_frames = len(image_paths)
        results = {
            'folder_name': folder_name,
            'frame_results': frame_results,
            'face_detection_rate': face_detection_count / total_frames if total_frames > 0 else 0.0,
            'emotion_summary': self._summarize_emotions(emotion_data)
        }
        
        print(f"子文件夹 '{folder_name}' 分析完成!")
        print(f"人脸检测率: {results['face_detection_rate']*100:.1f}%")
        
        return results
    
    def _summarize_emotions(self, emotions):
        """统计情绪分布"""
        if not emotions:
            return {}
        emotion_counts = {}
        for emotion in emotions:
            emotion_counts[emotion] = emotion_counts.get(emotion, 0) + 1
        total = len(emotions)
        return {emotion: count / total * 100 for emotion, count in emotion_counts.items()}

class ResultsGenerator:
    """负责生成和保存分析结果，仅输出HTML和CSV"""
    
    def __init__(self, results, output_base_folder=OUTPUT_FOLDER):
        self.results = results
        self.output_base_folder = output_base_folder
    
    def save_results(self, video_id=None):
        """保存HTML和CSV结果，使用 video_id 指定输出文件夹"""
        if video_id:
            output_folder = os.path.join(self.output_base_folder, video_id)
        else:
            output_folder = os.path.join(self.output_base_folder, self.results['folder_name'])
        
        os.makedirs(output_folder, exist_ok=True)
        print(f"正在将结果保存到 '{output_folder}' 文件夹...")
        
        self.save_csv_results(output_folder)
        self.generate_html_report(output_folder)
        print(f"结果已保存到 '{output_folder}' 文件夹")
    
    def save_csv_results(self, output_folder):
        """保存CSV格式的帧分析结果"""
        csv_file = os.path.join(output_folder, "frame_analysis.csv")
        with open(csv_file, 'w', newline='', encoding='utf-8') as f:
            writer = csv.writer(f)
            writer.writerow(['frame_path', 'eye_contact', 'expression', 'posture', 'face_dynamics', 'emotion'])
            for r in self.results['frame_results']:
                writer.writerow([
                    os.path.basename(r['path']), 
                    round(r['eye_contact'], 2),
                    round(r['expression'], 2),
                    round(r['posture'], 2),
                    round(r['face_dynamics'], 2),
                    r['emotion']
                ])
    
    def generate_html_report(self, output_folder):
        """生成HTML分析报告"""
        html_report = f"""<!DOCTYPE html>
<html>
<head>
    <meta charset="UTF-8">
    <title>{self.results['folder_name']} 图像帧分析报告</title>
    <style>
        body {{ font-family: Arial, sans-serif; margin: 20px; line-height: 1.6; }}
        h1, h2 {{ color: #2c3e50; }}
        table {{ border-collapse: collapse; width: 100%; margin: 20px 0; }}
        th, td {{ border: 1px solid #ddd; padding: 8px; text-align: left; }}
        th {{ background-color: #f2f2f2; }}
        tr:nth-child(even) {{ background-color: #f9f9f9; }}
    </style>
</head>
<body>
    <h1>{self.results['folder_name']} 图像帧分析报告</h1>
    
    <h2>分析概览</h2>
    <ul>
        <li><strong>总帧数:</strong> {len(self.results['frame_results'])}</li>
        <li><strong>人脸检测率:</strong> {self.results['face_detection_rate']*100:.1f}%</li>
    </ul>
    
    <h2>情绪分布</h2>
    <ul>
        {"".join([f"<li><strong>{emotion}:</strong> {percentage:.1f}%</li>" for emotion, percentage in self.results['emotion_summary'].items()])}
    </ul>
    
    <h2>帧分析详情</h2>
    <table>
        <tr>
            <th>帧</th>
            <th>眼神接触</th>
            <th>表情</th>
            <th>姿势</th>
            <th>面部动态</th>
            <th>情绪</th>
        </tr>
        {"".join([f"<tr><td>{os.path.basename(frame['path'])}</td><td>{frame['eye_contact']:.2f}</td><td>{frame['expression']:.2f}</td><td>{frame['posture']:.2f}</td><td>{frame['face_dynamics']:.2f}</td><td>{frame['emotion']}</td></tr>" for frame in self.results['frame_results']])}
    </table>
    
    <footer>
        <p>报告生成时间: {datetime.now().strftime("%Y-%m-%d %H:%M:%S")}</p>
    </footer>
</body>
</html>
"""
        with open(os.path.join(output_folder, 'analysis_report.html'), 'w', encoding='utf-8') as f:
            f.write(html_report)

def main(folder_path=None):
    """主函数 - 支持分析单个文件夹或整个frames文件夹"""
    print("=" * 50)
    print("图像帧视觉分析系统")
    print("=" * 50)
    
    analyzer = SimpleImageAnalyzer()
    
    if folder_path:  # 处理单个文件夹
        if not os.path.exists(folder_path):
            print(f"错误: 未找到'{folder_path}'文件夹")
            return None
        frame_count = len([f for f in os.listdir(folder_path) 
                         if f.lower().endswith(('.jpg', '.jpeg', '.png'))])
        if frame_count == 0:
            print(f"警告: 文件夹 '{folder_path}' 无图片可分析")
            return None
            
        print(f"图片帧数: {frame_count}")
        results = analyzer.analyze_image_folder(folder_path)
        if results:
            generator = ResultsGenerator(results)
            generator.save_results()
            print(f"文件夹 '{folder_path}' 处理完成")
        return results
    
    # 处理整个 FRAMES_FOLDER
    if not os.path.exists(FRAMES_FOLDER):
        print(f"错误: 未找到'{FRAMES_FOLDER}'文件夹")
        return 1
    
    subfolders = [os.path.join(FRAMES_FOLDER, f) 
                 for f in os.listdir(FRAMES_FOLDER) 
                 if os.path.isdir(os.path.join(FRAMES_FOLDER, f))]
    
    if not subfolders:
        print(f"错误: '{FRAMES_FOLDER}' 中没有找到子文件夹")
        return 1
    
    print(f"发现 {len(subfolders)} 个子文件夹待分析")
    
    for folder_path in subfolders:
        print("-" * 50)
        folder_name = os.path.basename(folder_path)
        print(f"开始处理子文件夹: {folder_name}")
        
        frame_count = len([f for f in os.listdir(folder_path) 
                         if f.lower().endswith(('.jpg', '.jpeg', '.png'))])
        if frame_count == 0:
            print(f"警告: 子文件夹 '{folder_name}' 无图片可分析")
            continue
            
        print(f"图片帧数: {frame_count}")
        
        try:
            results = analyzer.analyze_image_folder(folder_path)
            if results:
                generator = ResultsGenerator(results)
                generator.save_results()
                print(f"子文件夹 '{folder_name}' 处理完成")
            else:
                print(f"子文件夹 '{folder_name}' 分析失败")
        except Exception as e:
            print(f"处理子文件夹 '{folder_name}' 时出错: {str(e)}")
    
    print("=" * 50)
    print("所有子文件夹分析完成！")
    print(f"结果已保存到 '{OUTPUT_FOLDER}' 文件夹")
    return 0

if __name__ == "__main__":
    sys.exit(main())